package org.amos.limiter.components;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Collections;
import java.util.Objects;

/**
 * 基于redis+lua的令牌桶限流器
 *
 * @author admin
 */
@Slf4j
@Component
public class TokenBucketRateLimiter extends AbstractRateLimiter {

    @Autowired
    private DefaultRedisScript<Boolean> script;
    @Resource(name = "luaRedisTemplate")
    private RedisTemplate<String, Object> redisTemplate;

    /**
     * 限流检测(单个接口)
     *
     * @param name         需要限流的接口名
     * @param permits      最大令牌数
     * @param timeInterval 生成的令牌速率
     * @return 是否通过限流 true: 通过
     */
    @Override
    protected boolean acquire(String name, long permits, long timeInterval) {
        // 错误的参数将不起作用
        if (timeInterval <= 0 || timeInterval <= 0) {
            log.warn("maxPermits and tokensPerSeconds can not be less than zero...");
            return true;
        }
        long maxBurstTime = 500;
        long bucketMaxTokens = permits;
        // 限流时间间隔
        long resetBucketInterval = timeInterval;
        // 令牌的产生间隔 = 限流时间 / 最大令牌数
        long intervalPerPermit = resetBucketInterval / bucketMaxTokens;
        // 初始令牌数 = 最大的突发流量的持续时间 / 令牌产生间隔
        // 用 最大的突发流量的持续时间 计算的结果更加合理,并不是每次初始化都要将桶装满
        long initTokens = Math.min(maxBurstTime / intervalPerPermit, bucketMaxTokens);

        // 参数结构: KEYS = [限流的key]   ARGV = [最大令牌数, 生成的间隔(ms), 本次请求的毫秒数]
        // 参数1: 脚本, 参数2: 脚本中的KEYS数组, 参数3: 脚本中的ARGV数组
        Boolean result = this.redisTemplate.execute(this.script, Collections.singletonList(name), intervalPerPermit, System.currentTimeMillis(), initTokens, bucketMaxTokens, resetBucketInterval);
        log.error("RateLimiter,value:{}", result);
        return Objects.nonNull(result) && result;
    }

}
